#!/usr/bin/env python3
# -*- coding:utf-8 -*-

"""
【节点名称】：
    RTSPStreamerNode
【依赖项安装】：

    1 安装 GStreamer-1.0及相关插件:
        sudo apt-get install libgstreamer1.0-0 gstreamer1.0-plugins-base gstreamer1.0-plugins-good gstreamer1.0-plugins-bad gstreamer1.0-plugins-ugly gstreamer1.0-libav gstreamer1.0-doc gstreamer1.0-tools gstreamer1.0-x gstreamer1.0-alsa gstreamer1.0-gl gstreamer1.0-gtk3 gstreamer1.0-qt5 gstreamer1.0-pulseaudio
    2 安装 RTSP 服务器:
        sudo apt-get install libglib2.0-dev libgstrtspserver-1.0-dev gstreamer1.0-rtsp
    3 安装 OPENCV-PYTHON
        pip install opencv-python
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
    sensor_msgs::CameraCalibration (输入相机参数)
【发布类型】：
    无
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号,如果需要长期后台执行,建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    specified_output_topic (str): 指定输出的话题地址
【节点参数】：

【备注】：
    无
"""

import threading
import time
import cv2
import os
import json
from queue import Queue
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, BaseNode, get_extra_args
import argparse
from typing import Union

import subprocess
import gi
gi.require_version('Gst', '1.0')
gi.require_version('GstRtspServer', '1.0')
from gi.repository import Gst, GstRtspServer, GLib


class GStreamerNode_Cuda(threading.Thread, BaseNode):
    def __init__(self, job_name, ip='127.0.0.1', port=9094, param_dict_or_file=None, sms_shutdown=True, **kwargs):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )

        self.image_topic = self.get_param("image_topic", "/live/sensor/image_raw")
        self.resize = self.get_param("resize", [640, 480])
        self.fps = self.get_param("fps", 30)
        self.rtsp_port = self.get_param("port", 8554)
        self.stream_uri = self.get_param("stream_uri", "/live")
        self.params_help()

        self.queue = Queue()
        self.queue_pool.append(self.queue)

        self.subscriber = Subscriber(
            self.image_topic,
            'std_msgs::Null',
            self.image_callback,
            ip=ip,
            port=port
        )

        self.start()

        self.number_frames = 0
        self.duration = 1 / self.fps * Gst.SECOND

        # self.launch_string = f'appsrc name=source is-live=true block=true format=GST_FORMAT_TIME ' \
        #                      f'caps=video/x-raw,format=BGR,width={self.resize[0]},height={self.resize[1]},framerate={self.fps}/1 ' \
        #                      f'! videoconvert ! nvvidconv ! video/x-raw(memory:NVMM) ' \
        #                      f'! nvv4l2h264enc insert-sps-pps=true bitrate=2048000 ! h264parse ' \
        #                      f'! rtph264pay config-interval=1 name=pay0 pt=96'
        self.launch_string = 'appsrc name=source is-live=true block=true format=GST_FORMAT_TIME ' \
                        'caps=video/x-raw,format=BGR,width={},height={},framerate={}/1 ' \
                        '! videoconvert ! nvvidconv ! video/x-raw(memory:NVMM) ' \
                        '! nvv4l2h264enc  insert-sps-pps=true bitrate= 2048000 ! h264parse' \
                        '! rtph264pay config-interval=1 name=pay0 pt=96' \
                        .format(self.resize[0], self.resize[1], self.fps)

        self.init_gstreamer()

    def image_callback(self, msg):
        frame = sms2cvimg(msg)
        if self.queue.qsize() >= 100:
            try:
                self.queue.get(block=False)
            except:
                pass
        self.queue.put(frame)

    def init_gstreamer(self):
        Gst.init(None)

        factory = self.create_sensor_factory()
        server = GstRtspServer.RTSPServer()
        server.set_service(str(self.rtsp_port))
        server.get_mount_points().add_factory(self.stream_uri, factory)
        server.attach(None)

        self.loop = GLib.MainLoop()
        self.gst_thread = threading.Thread(target=self.loop.run)
        self.gst_thread.daemon = True
        self.gst_thread.start()

    def create_sensor_factory(self):
        node = self

        class SensorFactory(GstRtspServer.RTSPMediaFactory):
            def __init__(self):
                super().__init__()
                self.set_shared(True)
                self.launch_string = node.launch_string

            def do_create_element(self, url):
                return Gst.parse_launch(self.launch_string)

            def do_configure(self, rtsp_media):
                node.number_frames = 0
                appsrc = rtsp_media.get_element().get_child_by_name("source")
                appsrc.connect("need-data", node.on_need_data)

        return SensorFactory()

    def on_need_data(self, src, length):
        if not self.is_running:
            return
        frame = self.queue.get()
        while not self.queue.empty():
            frame = self.queue.get()

        frame = cv2.resize(frame, self.resize, interpolation=cv2.INTER_LINEAR)
        data = frame.tobytes()

        buf = Gst.Buffer.new_allocate(None, len(data), None)
        buf.fill(0, data)
        buf.duration = self.duration
        timestamp = self.number_frames * self.duration
        buf.pts = buf.dts = int(timestamp)
        buf.offset = timestamp
        self.number_frames += 1

        retval = src.emit('push-buffer', buf)
        print('pushed buffer, frame {}, duration {} ns, durations {} s'.format(self.number_frames,
                                                                            self.duration,
                                                                            self.duration / Gst.SECOND))
        if retval != Gst.FlowReturn.OK:
            print("Push buffer failed:", retval)

    def release(self):
        self.subscriber.kill()
        if self.loop:
            self.loop.quit()
        super().release()


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = GStreamerNode_Cuda(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()
